package com.tgy.tree;

import java.util.Comparator;

/**
 * @Author: tgy
 * @Date: 2020/9/14 9:59 下午
 */
public class AVLTree<E> extends BinarySearchTree<E> {

  public AVLTree() {

    this(null);
  }

  public AVLTree(Comparator<E> comparator) {
    super(comparator);
  }

  private enum RotateEnum{
    /**
     * 向左旋转
     */
    LEFT,
    /**
     * 向右旋转
     */
    RIGHT

  }

  private static class AVLNode<E> extends Node<E> {

    int height = 1;

    public AVLNode(E element, Node<E> parent) {
      super(element, parent);
    }

    /**
     * 获取左子树高度
     * @return
     */
    private int leftHeight() {

      return left != null ? ((AVLNode<E>)left).height:0;
    }

    /**
     * 获取右子树的高度
     * @return
     */
    private int rightHeight(){

      return right != null ? ((AVLNode<E>)right).height:0;
    }

    /**
     * 获取节点的平衡因子
     * @return
     */
    public int balanceFactor(){

      return leftHeight() - rightHeight();
    }

    public boolean isBalance() {

      return Math.abs(balanceFactor()) <= 1;
    }


    /**
     * 增加树高
     */
    public void addHeight() {

      height = 1 + Math.max(leftHeight(),rightHeight());
    }

    /**
     * 获取avl树在旋转过程中的pNode
     * @return
     */
    public AVLNode<E> getPNode() {

      if (leftHeight() > rightHeight()) {

        return leftAVLNode();
      }else {

        return rightAVLNode();
      }
    }


    /**
     *
     * @return
     */
    public AVLNode<E> leftAVLNode() {

      return (AVLNode<E>) left;
    }

    public AVLNode<E> rightAVLNode(){

      return (AVLNode<E>) right;
    }

    public AVLNode<E> parentAVLNode() {

      return (AVLNode<E>)parent;
    }

    @Override
    public String toString() {
      return element + "_h(" + height+")";
    }
  }

  /**
   * 创建AVLNode节点的树
   * @param element 节点元素
   * @param parentNode 父节点
   * @return
   */
  @Override
  protected Node<E> createNode(E element, Node<E> parentNode) {

    return new AVLNode<>(element,parentNode);
  }

  /**
   * 添加完元素之后对树进行调整
   * @param node
   */
  @Override
  protected void addAdjustTree(Node<E> node) {

    AVLNode<E> avlNode = (AVLNode<E>) node;
    avlNode = avlNode.parentAVLNode();
    while (true) {

      if (avlNode != null && !avlNode.isBalance()){

        doAdjuestTree(avlNode);
        break;
      }else {

        if (avlNode == null) {
          break;
        }
        avlNode.addHeight();
        avlNode = avlNode.parentAVLNode();
      }
    }
  }

  @Override
  protected void removeAdjustTree(Node<E> node){

    AVLNode<E> avlNode = (AVLNode<E>) node;
    avlNode = avlNode.parentAVLNode();

    while (true) {

      if (avlNode != null && !avlNode.isBalance()){

        doAdjuestTree(avlNode);
      }else {

        if (avlNode == null) {
          break;
        }
        avlNode.addHeight();
        avlNode = avlNode.parentAVLNode();
      }
    }

  }

  private void doAdjuestTree(AVLNode<E> gNode) {

    AVLNode<E> pNode = gNode.getPNode();

    if (pNode.equals(gNode.left)) {
//      左 L
      if (pNode.balanceFactor() > 0){
//        L
          rotate(gNode,RotateEnum.RIGHT);
      }else {
//        R
          rotate(pNode,RotateEnum.LEFT);
          rotate(gNode,RotateEnum.RIGHT);
      }

    }else {
//      R
      if (pNode.balanceFactor() > 0){
//        L
        rotate(pNode,RotateEnum.RIGHT);
        rotate(gNode,RotateEnum.LEFT);
      }else {
//        R
        rotate(gNode,RotateEnum.LEFT);
      }
    }
  }

  private void rotate(AVLNode<E> gNode,RotateEnum rotateEnum){

    AVLNode<E> pNode = rotateEnum == RotateEnum.LEFT? gNode.rightAVLNode():gNode.leftAVLNode();

    if (gNode.parent == null){

      root = pNode;
    }else {

      if (gNode.equals(gNode.parent.left)){

        gNode.parent.left = pNode;
      }else if (gNode.equals(gNode.parent.right)){

        gNode.parent.right = pNode;
      }
    }

    pNode.parent = gNode.parent;
    gNode.parent = pNode;

    if (rotateEnum == RotateEnum.LEFT) {

      if (pNode.left != null){
        pNode.left.parent = gNode;
      }

      gNode.right = pNode.left;
      pNode.left = gNode;
    }else {

      if (pNode.right != null){

        pNode.right.parent = gNode;
      }

      gNode.left  = pNode.right;
      pNode.right = gNode;
    }

    gNode.addHeight();
    pNode.addHeight();
  }

}
